import subprocess

import sys

class Logger(object):
    def __init__(self, filename='default.log', stream=sys.stdout):
        self.terminal = stream
        self.log = open(filename, 'a')

    def write(self, message):
        self.terminal.write(message)
        self.log.write(message)
        self.flush()

    def flush(self):
        self.terminal.flush()
        self.log.flush()

    def close(self):
        self.log.close()

sys.stdout = Logger('a.log', sys.stdout)
sys.stderr = Logger('a.log', sys.stderr)

# 定义要运行的Python脚本和参数
script = '../benchmark/run.py'
args = '-d cpu -t eval'
models = ['BERT_pytorch',
          'LearningToPaint',
          'Super_SloMo',
          'alexnet',
          'basic_gnn_edgecnn',
          'basic_gnn_gcn',
          'basic_gnn_gin',
          'basic_gnn_sage',
          'cm3leon_generate',
          'dcgan',
          'demucs',
          'densenet121',
          'detectron2_fasterrcnn_r_101_c4',
          'detectron2_fasterrcnn_r_101_dc5',
          'detectron2_fasterrcnn_r_101_fpn',
          'detectron2_fasterrcnn_r_50_c4',
          'detectron2_fasterrcnn_r_50_dc5',
          'detectron2_fasterrcnn_r_50_fpn',
          'detectron2_fcos_r_50_fpn',
          'detectron2_maskrcnn',
          'detectron2_maskrcnn_r_101_c4',
          'detectron2_maskrcnn_r_101_fpn',
          'detectron2_maskrcnn_r_50_c4',
          'detectron2_maskrcnn_r_50_fpn',
          'dlrm',
          'doctr_det_predictor',
          'doctr_reco_predictor',
          'drq',
          'fastNLP_Bert',
          'functorch_dp_cifar10',
          'functorch_maml_omniglot',
          'hf_Albert',
          'hf_Bart',
          'hf_Bert',
          'hf_Bert_large',
          'hf_BigBird',
          'hf_DistilBert',
          'hf_GPT2',
          'hf_GPT2_large',
          'hf_Longformer',
          'hf_Reformer',
          'hf_T5',
          'hf_T5_base',
          'hf_T5_generate',
          'hf_T5_large',
          'hf_Whisper',
          'hf_clip',
          'hf_distil_whisper',
          'lennard_jones',
          'llama',
          'llama_v2_7b_16h',
          'llava',
          'maml',
          'maml_omniglot',
          'microbench_unbacked_tolist_sum',
          'mnasnet1_0',
          'mobilenet_v2',
          'mobilenet_v2_quantized_qat',
          'mobilenet_v3_large',
          'moondream',
          'nanogpt',
          'opacus_cifar10',
          'phlippe_densenet',
          'phlippe_resnet',
          'pyhpc_equation_of_state',
          'pyhpc_isoneutral_mixing',
          'pyhpc_turbulent_kinetic_energy',
          'pytorch_CycleGAN_and_pix2pix',
          'pytorch_stargan',
          'pytorch_unet',
          'resnet152',
          'resnet18',
          'resnet50',
          'resnet50_quantized_qat',
          'resnext50_32x4d',
          'sam',
          'sam_fast',
          'shufflenet_v2_x1_0',
          'soft_actor_critic',
          'squeezenet1_1',
          'stable_diffusion_text_encoder',
          'stable_diffusion_unet',
          'timm_efficientnet',
          'timm_nfnet',
          'timm_regnet',
          'timm_resnest',
          'timm_vision_transformer',
          'timm_vision_transformer_large',
          'timm_vovnet',
          'torch_multimodal_clip',
          'tts_angular',
          'vgg16',
          'vision_maskrcnn',
          'yolov3'
         ]

try:
    for model in models:

        print('------------------------- ' + 'model: ' + model + ' -------------------------')

        # 构建命令字符串
        command = './perf_information_get.sh' + f' "python3 {script} {model} {args}"' + ' ./result'

        try:
            # 使用 Popen 运行命令，并捕获其输出
            process = subprocess.Popen(command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
            print()

            # 实时捕获并处理输出
            for line in iter(process.stdout.readline, b''):
                sys.stdout.write(line.decode())
            for line in iter(process.stderr.readline, b''):
                sys.stderr.write(line.decode())

            process.stdout.close()
            process.stderr.close()
            process.wait()

            if process.returncode != 0:
                print(f'Command failed with {process.returncode}')
                break
        except subprocess.CalledProcessError as e:
            print(f'Command failed with {e.returncode}')
            # 注释后，即使有依赖未安装也不会停止
            # break
        except KeyboardInterrupt:
            print('Process interrupted by user')
            process.terminate()
            break
finally:
    # 关闭日志文件
    sys.stdout.close()
    sys.stderr.close()
    # 恢复原始的 sys.stdout
    sys.stdout = sys.__stdout__
    sys.stderr = sys.__stderr__